from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta

fake = Faker('zh_CN')

# ========== 药品数据配置 ==========
ANTI_CANCER_DRUGS = [
    {
        'generic_name': '紫杉醇',
        'trade_names': ['泰素', '紫素', '安素泰'],
        'manufacs': ['百时美施贵宝', '恒瑞医药', '齐鲁制药'],
        'approval_no': '国药准字H{}'.format(random.randint(1000000, 9999999)),
        'forms': ['注射液', '脂质体'],
        'specs': ['30mg/5ml', '100mg/16.7ml'],
        'packing': ['1支/盒', '5支/盒']
    },
    {
        'generic_name': '顺铂',
        'trade_names': ['顺铂注射液', '诺欣'],
        'manufacs': ['齐鲁制药', '豪森药业'],
        'approval_no': '国药准字H{}'.format(random.randint(1000000, 9999999)),
        'forms': ['注射液'],
        'specs': ['10mg/支', '20mg/支'],
        'packing': ['1支/盒', '10支/盒']
    },
    {
        'generic_name': '卡铂',
        'trade_names': ['波贝', '卡铂注射液'],
        'manufacs': ['恒瑞医药', '正大天晴'],
        'approval_no': '国药准字H{}'.format(random.randint(1000000, 9999999)),
        'forms': ['注射液'],
        'specs': ['100mg/10ml'],
        'packing': ['1支/盒']
    },
    {
        'generic_name': '奥沙利铂',
        'trade_names': ['乐沙定', '艾恒'],
        'manufacs': ['赛诺菲', '恒瑞医药'],
        'approval_no': '国药准字H{}'.format(random.randint(1000000, 9999999)),
        'forms': ['注射液'],
        'specs': ['50mg/支', '100mg/支'],
        'packing': ['1支/盒', '5支/盒']
    },
    {
        'generic_name': '多西他赛',
        'trade_names': ['泰索帝', '多帕菲'],
        'manufacs': ['赛诺菲', '恒瑞医药'],
        'approval_no': '国药准字H{}'.format(random.randint(1000000, 9999999)),
        'forms': ['注射液'],
        'specs': ['20mg/支', '80mg/支'],
        'packing': ['1支/盒']
    },
    {
        'generic_name': '吉西他滨',
        'trade_names': ['健择', '泽菲'],
        'manufacs': ['礼来', '恒瑞医药'],
        'approval_no': '国药准字H{}'.format(random.randint(1000000, 9999999)),
        'forms': ['冻干粉针'],
        'specs': ['200mg/支', '1g/支'],
        'packing': ['1支/盒']
    },
    {
        'generic_name': '培美曲塞',
        'trade_names': ['力比泰', '普来乐'],
        'manufacs': ['礼来', '豪森药业'],
        'approval_no': '国药准字H{}'.format(random.randint(1000000, 9999999)),
        'forms': ['冻干粉针'],
        'specs': ['100mg/支', '500mg/支'],
        'packing': ['1支/盒']
    }
]

INSURANCE_TYPES = ['甲类医保', '乙类医保', '自费', '特殊药品']
HOSPITAL_DRUG_CODE_PREFIX = 'DRG'


# ========== 批次时间配置 ==========
def generate_month_range(start_month, end_month):
    """
    要求 start_month 和 end_month 必须是 YYYY-MM 这种格式
    :param start_month:
    :param end_month:
    :return:
    """
    """生成月份序列（格式：YYYYMM）"""
    # 字符串转datetime类型的
    start = datetime.strptime(start_month, "%Y-%m")
    print('-----start', start)  # 2024-08-01 00:00:00
    end = datetime.strptime(end_month, "%Y-%m")
    months = []

    current = start
    while current <= end:
        months.append(current.strftime("%Y%m"))
        # 把一个 datetime 对象 格式化成 字符串。
        # strftime 的意思是：string format time（把时间格式化成字符串）
        current += relativedelta(months=1)  # 每次加一个月，循环推进

    return months  # ['202301', '202302', '202303', '202304', '202305', '202306']


# 配置月份范围（可修改）
MONTH_RANGE = generate_month_range("2024-08", "2025-08")  # 多少个月
RECORDS_PER_MONTH = 15  # 每月生成15条记录（15 * 13=195≈200）   # 多少条


# ========== 核心生成函数 ==========
def generate_drug_purchase_records():
    """生成抗肿瘤药物采购记录"""
    sql_statements = []
    record_count = 0

    for month in MONTH_RANGE:  # 12个月
        for i in range(1, RECORDS_PER_MONTH + 1):
            if record_count >= 200:
                break

            record_count += 1
            # 开始赋值
            drug = random.choice(ANTI_CANCER_DRUGS)
            trade_name = random.choice(drug['trade_names'])
            manufac = random.choice(drug['manufacs'])
            form = random.choice(drug['forms'])
            spec = random.choice(drug['specs'])
            packing = random.choice(drug['packing'])

            # 生成药品价格（基于规格）
            base_price = {
                '紫杉醇': 800,
                '顺铂': 50,
                '卡铂': 200,
                '奥沙利铂': 600,
                '多西他赛': 1000,
                '吉西他滨': 300,
                '培美曲塞': 1500
            }.get(drug['generic_name'], 500)

            # 根据规格调整价格
            if '100mg' in spec:
                price = base_price * 2
            elif '50mg' in spec:
                price = base_price * 1.5
            else:
                price = base_price

            price = round(price * random.uniform(0.8, 1.2), 2)

            # 生成采购数据
            pur_quantity = random.randint(10, 100)
            pur_total_price = round(price * pur_quantity, 2)

            # 最后整合成一个json
            data = {
                'pur_sn': f"PUR{month}{i:03d}",
                'hos_drug_code': f"{HOSPITAL_DRUG_CODE_PREFIX}{random.randint(1000, 9999)}",
                'generic_name': drug['generic_name'],
                'trade_name': trade_name,
                'manufac': manufac,
                'approval_no': drug['approval_no'],
                'form': form,
                'spec': spec,
                'packing_quantity': packing.split('/')[0],
                'packing_unit': packing.split('/')[1],
                'pur_unit_price': price,
                'pur_quantity': pur_quantity,
                'pur_total_price': pur_total_price,
                'pur_date': fake.date_time_between(
                    start_date=datetime.strptime(month, "%Y%m"),
                    end_date=(datetime.strptime(month, "%Y%m") + relativedelta(months=1))
                ).strftime('%Y-%m-%d %H:%M:%S'),
                'insurance_type': random.choice(INSURANCE_TYPES),
                'extend_data1': f"批号：{random.randint(100000, 999999)}",
                'extend_data2': f"有效期至：{(datetime.now() + timedelta(days=random.randint(180, 720))).strftime('%Y-%m-%d')}",
                'yy_upload_status': 0,
                'yy_etl_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'yy_upload_time': None,
                'yy_batch_time': month,
                'yy_record_batch_id': f"BATCH{month}_{i}",
                'yy_backfill_time': None,
                'yy_backfill_status': None
            }

            sql = _generate_sql('b01_1', data)
            print(sql)
            sql_statements.append(sql)
            """"
            INSERT INTO `b01_1` (`pur_sn`, `hos_drug_code`, `generic_name`, `trade_name`, `manufac`, `approval_no`, `form`, `spec`, `packing_quantity`, `packing_unit`, `pur_unit_price`, `pur_quantity`, `pur_total_price`, `pur_date`, `insurance_type`, 
            `extend_data1`, `extend_data2`, `yy_upload_status`, `yy_etl_time`, 
            `yy_upload_time`, `yy_batch_time`, `yy_record_batch_id`, `yy_backfill_time`, 
            `yy_backfill_status`) VALUES ('PUR202508015', 'DRG1977', '顺铂', '诺欣', '豪森药业', '国药准字H8106408', '注射液', '10mg/支', '1支', '盒', 56.44, 94, 5305.36, '2025-08-09 17:36:11', '乙类医保', '批号：223573', '有效期至：2027-03-27', 0, '2025-10-03 19:55:34', NULL, '202508', 'BATCH202508_15', NULL, NULL);
            """

    return sql_statements


def _generate_sql(table, data):
    """
    传入表和data
    :param table:
    :param data:
    :return:
    字典值
    """
    """生成SQL插入语句"""
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    # print('--------',type(columns))  str类型的
    # `patient_id`, `visit_sn`, `name`
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')  # 如果值是 None（Python 里的空值），SQL 里要用 NULL 表示，不能写 'None'。
        # elif isinstance(v, (int, float)):
        #     # values.append(str(v))
        #     pass
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")

    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)})"

    # v = "张三"：
    #
    # str(v) → "张三"
    #
    # 包装成 f"'{escaped_value}'" → "'张三'"
    #
    # 拼到 SQL：
    #
    # VALUES ('张三')


#     str(v) 是把 Python 的值转成字符串，方便拼接 SQL。
#
# 是否加 '...' 是根据 SQL 类型规则决定的：
#
# 数字：只用 str(v)，不加 '...'
#
# 字符串：必须 '...' 包裹
#
# None：写成 NULL


# ========== 执行生成 ==========
if __name__ == "__main__":
    print("开始生成抗肿瘤药物采购记录(b01_1表)...")
    print(f"月份范围：{MONTH_RANGE[0]} 至 {MONTH_RANGE[-1]}")
    print(f"共{len(MONTH_RANGE)}个月份，每月{RECORDS_PER_MONTH}条记录")

    records = generate_drug_purchase_records()

    with open('b01_1_drug_purchases.sql', 'w', encoding='utf-8') as f:
        f.write("-- 抗肿瘤药物采购记录数据(b01_1表)\n")
        f.write(f"-- 生成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"-- 月份范围：{MONTH_RANGE[0]} 至 {MONTH_RANGE[-1]}\n")
        f.write(f"-- 共{len(records)}条记录\n\n")

        # 分批写入防止内存溢出
        for i in range(0, len(records), 50):
            batch = records[i:i + 50]
            f.write("\n".join(batch) + "\n")
            print(f"已生成 {i + 50 if i + 50 < len(records) else len(records)} 条记录")

    print("\n生成完成！文件已保存为 b01_1_drug_purchases.sql")
